from app.strategy.base import BaseStrategy
from app.gnn.infer import gnn_signal
import pandas as pd
from loguru import logger
class GridGNNStrategy(BaseStrategy):
    
    def __init__(self, codes, start, end, **kwargs):
        # 如果父类还需要别的参数，继续往 **kwargs 里传
        super().__init__(codes, start, end, **kwargs)
        
    def next(self, row: pd.Series) -> dict[str, float]:
        score = gnn_signal()                     # 当前 GNN 分数
        if not score:
            result = {c: 1.0 / len(self.codes) for c in self.codes}            
            logger.debug(f"GridGNNStrategy.next: {result.bars.columns.tolist()}")
            print(f"GridGNNStrategy.next: {result.bars.columns.tolist()}")
            logger.debug(f"GridGNNStrategy.next: {result.bars.head()}")
            print(f"GridGNNStrategy.next: {result.bars.head()}")
            return result
        total = sum(score.values())
        result = {c: score.get(c, 0) / total for c in self.codes}
        logger.debug(f"GridGNNStrategy.next: {result.bars.columns.tolist()}")
        print(f"GridGNNStrategy.next: {result.bars.columns.tolist()}")
        logger.debug(f"GridGNNStrategy.next: {result.bars.head()}")
        print(f"GridGNNStrategy.next: {result.bars.head()}")
        return result
